Skip to content

Update RefModel to use Episode #139

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 7 commits into from
Sep 15, 2025
Merged

Update RefModel to use Episode #139

merged 7 commits into from
Sep 15, 2025

Conversation

Jack-Khuu
Copy link
Contributor

@Jack-Khuu Jack-Khuu commented Sep 8, 2025 •
edited
Loading

Mini update to (titan) ReferenceModel signature to match RefModel of apps/GRPO/main by using Episode as input

Git refuses to pick up the file rename (reference_actor.py => reference_model.py), so I'll callout explicitly what was changed in reference_model.py

  • Deleted code under the experimental section
  • Rename TitanRefModel => ReferenceModel
  • Updated forward to use Episode (which wraps req/resp)
async def forward(self, episode: 'Episode') -> torch.Tensor:
async def forward(self, request: list[int], response: list[int]) -> torch.Tensor:

python apps/grpo/main.py

@meta-cla meta-cla bot added the CLA Signed This label is managed by the Meta Open Source bot. label Sep 8, 2025
pbontrager
pbontrager approved these changes Sep 15, 2025
View reviewed changes
Copy link
Contributor

@pbontrager pbontrager left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good. I'd just request that you import the logprobs computation from the one in trainer when it lands.

Also something to think about would be a way to reuse the trainer user defined batching logic here instead of using Episode.

@Jack-Khuu
Copy link
Contributor Author

Jack-Khuu commented Sep 15, 2025 •
edited
Loading

import the logprobs computation from the one in trainer when it lands.

trainer.py doesn't have a logprobs calc, though the one in rl/main.py Is functionally equivalent to the one here. I'll have a separate PR combining the 3 (rl/grpo/ref)

@Jack-Khuu Jack-Khuu merged commit 0b585c0 into main Sep 15, 2025
5 checks passed
@Jack-Khuu Jack-Khuu deleted the update-refmodel-episode branch September 15, 2025 22:57
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@pbontrager pbontrager pbontrager approved these changes

@allenwang28 allenwang28 Awaiting requested review from allenwang28

@joecummings joecummings Awaiting requested review from joecummings

Assignees

No one assigned

Labels

CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@Jack-Khuu @pbontrager